import PyPDF2
from reportlab.lib.pagesizes import letter
from reportlab.lib.units import inch
from reportlab.pdfgen import canvas
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.pdfbase import pdfmetrics
from datetime import datetime

from typing import List
from pikepdf import Pdf,Page,Rectangle


'''
向目标pdf文件批量添加水印
target_pdf_path:目标pdf文件路径+文件名
watermark_pad_path:水印pdf文件路径+文件名
nrow:水印平铺的行数
ncol:水印平铺的列数
skip_pages:需要跳过不添加水印的页数

'''
class Watermark:
    def __init__(self, 
                skip_pages:List[int] = [],
            ) -> None:        
        """
        为PDF文件添加水印。

        参数:
        watermark_pdf_path (str): 水印PDF文件的路径。
        nrow (int): 每行添加的水印数量。
        ncol (int): 每列添加的水印数量。
        skip_pages (List[int], 可选): 需要跳过不添加水印的页面编号列表，默认为空列表。

        返回:
        None
        """
        self.watermark_pdf_path = 'watermark.pdf'
        self.nrow = 1
        self.ncol = 1
        self.skip_pages = skip_pages

    def  generator_score_watermark(self, teacher, score, center_x=100, center_y=100) -> None:
        # Register Chinese font
        pdfmetrics.registerFont(TTFont('song', 'simsun.ttc'))

        c = canvas.Canvas(self.watermark_pdf_path, pagesize=(200,200))
        
        # Draw watermark circle
        c.setLineWidth(width=4)
        c.setStrokeColorRGB(0.1, 0.6, 0.8)
        c.setFillColorRGB(0.1, 0.6, 0.8)
        c.circle(x_cen=center_x, y_cen=center_y, r=60)
        
        c.setLineWidth(width=1)
        c.setStrokeColorRGB(0.1, 0.6, 0.8)
        c.setFillColorRGB(0.1, 0.6, 0.8)
        c.circle(x_cen=center_x, y_cen=center_y, r=55)

        # Add watermark text in Chinese
        c.setFont("song", 10)  # Use the registered Chinese font
        c.drawString(center_x-20, center_y+10, '成绩：{}'.format(str(score)))
        c.drawString(center_x-30, center_y-10, '阅卷人：{}'.format(teacher))
        c.drawString(center_x-20, center_y+30, '★ ★ ★')

        # Add timestamp
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        c.setFont("song", 9)
        c.drawString(center_x-40, center_y-30, f"{timestamp}")

        c.save()
    def add_score_watermark(self, pdf_path ) -> None:

        
        #打开目标PDF文件
        target_pdf = Pdf.open(pdf_path, allow_overwriting_input=True)

        #打开水印PDF文件并提取第一页作为水印
        watermark_pdf = Pdf.open(self.watermark_pdf_path)
        watermark_page = watermark_pdf.pages[0]

        #遍历目标PDF的每一页，根据行列数批量添加水印
        for idx,target_page in enumerate(target_pdf.pages):
            if idx in self.skip_pages:  # 如果当前页码在skip_pages列表中，则跳过
                continue
            for x in range(self.ncol):
                for y in range(self.nrow):
                    #计算水印在目标页面上的位置并添加
                    target_page.add_overlay(
                        watermark_page,
                        Rectangle(
                            # 右上角水印左下角X坐标，通常设为距离页面右侧边缘的小距离（如10单位）
                            target_page.trimbox[2] - watermark_page.trimbox[2] ,
                            # 右上角水印左下角Y坐标，通常设为距离页面顶部边缘的小距离（如10单位）
                            target_page.trimbox[3] - watermark_page.trimbox[3] ,
                            # 右上角水印右上角X坐标，等于左下角X坐标加上水印宽度
                            target_page.trimbox[2] ,
                            # 右上角水印右上角Y坐标，等于左下角Y坐标加上水印高度
                            target_page.trimbox[3] 
                        )
                    )
        #保存添加了水印的PDF文件，并在原文件名基础上添加"已添加水印"的后缀
        target_pdf.save(pdf_path)



    def make_score_watermark(self,pdf_path, teacher,score ):
        self.generator_score_watermark(teacher,score)
        self.add_score_watermark(pdf_path)
        import os 
        os.remove(self.watermark_pdf_path)


if __name__ == '__main__':
    water = Watermark()
    water.make_score_watermark('/nfs/bigdata2024/7/homework/5/stuSubmit/20204249/Borg, Omega, and kubernetes.pdf','张三',90)